/*
 * Copyright (C) 2010 Imavis Srl
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as published
 * by the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
package com.imavis.avimuxer;

import java.io.ByteArrayOutputStream;
import java.io.DataOutput;
import java.io.DataOutputStream;
import java.io.File;
import java.io.IOException;
import java.io.RandomAccessFile;
import java.util.Date;

/**
 * The AVI muxer.
 * <p>
 * Handles the RIFF AVI format and wraps the media data into an AVI container.
 *
 * @author Andrea Gualano
 */
final class AviMuxerImpl implements AviMuxer {
    private static final FourCC RIFF = new FourCC("RIFF");
    private static final FourCC AVI = new FourCC("AVI ");
    private static final FourCC LIST = new FourCC("LIST");
    private static final FourCC HDRL = new FourCC("hdrl");
    private static final FourCC AVIH = new FourCC("avih");
    private static final FourCC STRL = new FourCC("strl");
    private static final FourCC STRH = new FourCC("strh");
    private static final FourCC VIDS = new FourCC("vids");
    private static final FourCC STRF = new FourCC("strf");
    private static final FourCC MOVI = new FourCC("movi");
    private static final FourCC IDX1 = new FourCC("idx1");
    private static final FourCC STREAM_00DC = new FourCC("00dc");
    private int rate;
    private int scale;
    private int width;
    private int height;
    private final FourCC codec;
    private final RandomAccessFile out;
    private final ByteArrayOutputStream index;
    private final DataOutput indexDataOutput;
    private long riffOffset;
    private long moviOffset;
    private long aviMainHeaderOffset;
    private long aviStreamHeaderOffset;
    private long bitmapInfoHeaderOffset;
    private int frames;
    private long length;

    AviMuxerImpl(int rate, int scale, int width, int height, FourCC codec,
            File file) throws IOException {
        this.rate = rate;
        this.scale = scale;
        this.width = width;
        this.height = height;
        this.codec = codec;
        this.out = new RandomAccessFile(file, "rw");
        this.index = new ByteArrayOutputStream();
        this.indexDataOutput = new DataOutputStream(index);
    }

    /**
     * Writes all the headers to the output file.
     */
    @Override
    public void open() throws IOException {
        riffOffset = startList(RIFF, AVI);
        long hdrlOffset = startList(LIST, HDRL);
        writeAviMainHeader();
        long strlOffset = startList(LIST, STRL);
        writeAviStreamHeader();
        writeBitmapInfoHeader();
        endList(strlOffset);
        endList(hdrlOffset);
        moviOffset = startList(LIST, MOVI);
        length = out.getFilePointer();
    }

    /**
     * Adds a frame to the first video stream ("00dc").
     */
    @Override
    public void write(byte[] data, Date timestamp, boolean isKeyFrame)
            throws IOException {
        addToIndex(STREAM_00DC, isKeyFrame, length, data.length);
        writeChunk(STREAM_00DC, data);
        frames++;
        length = out.getFilePointer();
    }

    /**
     * Sets the framerate.
     * The new value is written to the file immediately.
     *
     * @param rate frame rate numerator (in Hz)
     * @param scale frame rate denominator
     * @see #close()
     */
    @Override
    public void setFramerate(int rate, int scale) throws IOException {
        this.rate = rate;
        this.scale = scale;
        long current = out.getFilePointer();
        writeAviMainHeader();
        writeAviStreamHeader();
        out.seek(current);
    }

    /**
     * Sets the resolution.
     * The new value is written to the file immediately.
     *
     * @param width horizontal resolution
     * @param height vertical resolution
     * @see #close()
     */
    @Override
    public void setResolution(int width, int height) throws IOException {
        if ((width != this.width) || (height != this.height)) {
            this.width = width;
            this.height = height;
            long current = out.getFilePointer();
            writeAviMainHeader();
            writeAviStreamHeader();
            writeBitmapInfoHeader();
            out.seek(current);
        }
    }

    /**
     * Writes end-of-file values to the headers, then closes the output file.
     */
    @Override
    public void close() throws IOException {
        endList(moviOffset);
        writeOldAviIndex();
        endList(riffOffset);
        writeAviMainHeader();
        writeAviStreamHeader();
        out.close();
    }

    /**
     * Returns the predicted length of the output file, if closed now.
     */
    @Override
    public long length() {
        return length + 8 + 16 * frames;
    }

    /**
     * Writes a 32-bit 4CC value to the file.
     */
    private void writeFourCC(FourCC fcc) throws IOException {
        out.writeInt(fcc.toInt());
    }

    /**
     * Writes a little-endian 32-bit integer to the file.
     */
    private void writeInteger(int i) throws IOException {
        out.writeInt(Integer.reverseBytes(i));
    }

    /**
     * Writes a little-endian 16-bit integer to the file.
     */
    private void writeShort(int i) throws IOException {
        out.writeShort(Short.reverseBytes((short) i));
    }

    /**
     * Writes a non-list chunk structure to the file.
     */
    private void writeChunk(FourCC type, byte[] data) throws IOException {
        writeFourCC(type);
        writeInteger(data.length);
        out.write(data);
        if (data.length % 2 == 1) {
            out.writeByte(0);
        }
    }

    /**
     * Writes a list header to the file.
     * The list must be closed later with {@link #endList(long)}
     *
     * @param header a list header ({@link #RIFF} or {@link #LIST})
     * @param type the actual list 4CC
     * @return the list offset (to be used in {@link #endList})
     */
    private long startList(FourCC header, FourCC type) throws IOException {
        writeFourCC(header);
        writeInteger(-1); // fill later
        long offset = out.getFilePointer();
        writeFourCC(type);
        return offset;
    }

    /**
     * Closes a list structure, writing the current length in the list header.
     * @param offset the list offset, returned by {@link #startList}
     */
    private void endList(long offset) throws IOException {
        long current = out.getFilePointer();
        out.seek(offset - 4);
        writeInteger((int) (current - offset));
        out.seek(current);
    }

    /**
     * Writes a "avih" chunk containing a AVIMAINHEADER structure.
     */
    private void writeAviMainHeader() throws IOException {
        if (aviMainHeaderOffset == 0) {
            aviMainHeaderOffset = out.getFilePointer();
        } else {
            out.seek(aviMainHeaderOffset);
        }
        writeFourCC(AVIH);
        writeInteger(56);
        // DWORD  dwMicroSecPerFrame;
        writeInteger((int) (1000000L * scale / rate));
        // DWORD  dwMaxBytesPerSec;
        writeInteger(0);
        // DWORD  dwPaddingGranularity;
        writeInteger(0);
        // DWORD  dwFlags;
        writeInteger(0x00000110); // AVIF_HASINDEX + AVIF_ISINTERLEAVED
        // DWORD  dwTotalFrames;
        writeInteger(frames);
        // DWORD  dwInitialFrames;
        writeInteger(0);
        // DWORD  dwStreams;
        writeInteger(1);
        // DWORD  dwSuggestedBufferSize;
        writeInteger(1024 * 1024);
        // DWORD  dwWidth;
        writeInteger(width);
        // DWORD  dwHeight;
        writeInteger(height);
        // DWORD  dwReserved[4];
        writeInteger(0);
        writeInteger(0);
        writeInteger(0);
        writeInteger(0);
    }

    /**
     * Writes a "strh" chunk containing a AVISTREAMHEADER structure.
     */
    private void writeAviStreamHeader() throws IOException {
        if (aviStreamHeaderOffset == 0) {
            aviStreamHeaderOffset = out.getFilePointer();
        } else {
            out.seek(aviStreamHeaderOffset);
        }
        writeFourCC(STRH);
        writeInteger(56);
        // FOURCC fccType;
        writeFourCC(VIDS);
        // FOURCC fccHandler;
        writeFourCC(codec);
        // DWORD  dwFlags;
        writeInteger(0);
        // WORD   wPriority;
        // WORD   wLanguage;
        writeInteger(0);
        // DWORD  dwInitialFrames;
        writeInteger(0);
        // DWORD  dwScale;
        writeInteger(scale);
        // DWORD  dwRate;
        writeInteger(rate);
        // DWORD  dwStart;
        writeInteger(0);
        // DWORD  dwLength;
        writeInteger(frames);
        // DWORD  dwSuggestedBufferSize;
        writeInteger(1024 * 1024);
        // DWORD  dwQuality;
        writeInteger(-1);
        // DWORD  dwSampleSize;
        writeInteger(0);
        // struct rcFrame;
        writeInteger(0);
        writeShort(width);
        writeShort(height);
    }

    /**
     * Writes a "strf" chunk containing a BITMAPINFOHEADER structure.
     */
    private void writeBitmapInfoHeader() throws IOException {
        if (bitmapInfoHeaderOffset == 0) {
            bitmapInfoHeaderOffset = out.getFilePointer();
        } else {
            out.seek(bitmapInfoHeaderOffset);
        }
        writeFourCC(STRF);
        writeInteger(40);
        // DWORD biSize;
        writeInteger(40);
        // LONG  biWidth;
        writeInteger(width);
        // LONG  biHeight;
        writeInteger(height);
        // WORD  biPlanes;
        writeShort(1);
        // WORD  biBitCount;
        writeShort(24);
        // DWORD biCompression;
        writeFourCC(codec);
        // DWORD biSizeImage;
        writeInteger(width * height * 3);
        // LONG  biXPelsPerMeter;
        writeInteger(0);
        // LONG  biYPelsPerMeter;
        writeInteger(0);
        // DWORD biClrUsed;
        writeInteger(0);
        // DWORD biClrImportant;
        writeInteger(0);
    }

    /**
     * Caches an index entry.
     */
    private void addToIndex(FourCC streamId, boolean isKeyFrame,
            long offset, int length) throws IOException {
        indexDataOutput.writeInt(streamId.toInt());
        indexDataOutput.writeInt(isKeyFrame ? 0x10000000 : 0);
        int o = (int) (offset - moviOffset);
        indexDataOutput.writeInt(Integer.reverseBytes(o));
        indexDataOutput.writeInt(Integer.reverseBytes(length));
    }

    /**
     * Writes a "idx1" chunk containing a an old-style (AVI 1.0) index.
     */
    private void writeOldAviIndex() throws IOException {
        writeChunk(IDX1, index.toByteArray());
    }
}
